import rospy
from std_srvs.srv import Empty
from sensor_msgs.msg import LaserScan
import math
import numpy as np
import json

RESOLUTION = 2*np.pi/3540

class Rplidar_Subscriber():
    INSTALL_ANGLE_BIAS = 315
    def __init__(self):
        self._scan = None

        rospy.Subscriber(
            "/scan",
            LaserScan,
            self.callback
        )

    def callback(self, data: LaserScan):
        self._scan = data

    @property
    def ladar_data(self):
        '''get fixed scan data'''
        if(self._scan):
            return self._fix(self._scan.ranges, RESOLUTION, 315) * 100
        else:
            return None

    def _fix(self, ranges, inc_rad, offset = INSTALL_ANGLE_BIAS):
        index = int(offset * math.pi / 180 / inc_rad)
        range_fixed = ranges[index:] + ranges[:index]
        return np.array(range_fixed)

    def start_motor(self):
        rospy.wait_for_service('start_motor')
        s = rospy.ServiceProxy('start_motor', Empty)
        s.call()
        s.close()

    def stop_motor(self):
        rospy.wait_for_service('stop_motor')
        s = rospy.ServiceProxy('stop_motor', Empty)
        s.call()
        s.close()


def range_filter(ranges, min = 100, max = 300, value = 0):
    ranges[ranges < min] = value
    ranges[ranges > max] = value


def range_filter_rect(ranges, x_max = 400, y_max = 300, min = 30, value = 0):
    '''warning: unit: m'''
    index_sep =  math.atan(y_max / x_max) / RESOLUTION
    length = len(ranges)

    # ds = []
    for i in range(length):
        if i < index_sep:
            d = x_max / math.cos( i * RESOLUTION )
            # ds.append(d)
            if ranges[i] > d or ranges[i] < min:
                ranges[i] = value
        
        else:
            d = y_max / math.cos( math.pi/2 - i * RESOLUTION )
            # ds.append(d)
            if ranges[i] > d or ranges[i] < min:
                ranges[i] = value
        
    # print(ds)

def angle_filter(ranges, the0, the1, value = 0):
    id0 = int(the0 / 180 * np.pi / RESOLUTION)
    id1 = int(the1 / 180 * np.pi / RESOLUTION)
    ranges[:id0] = value
    ranges[id1:] = value


# 找到杆子(最近的点)
def find_pole(ranges_valid):
    '''(angle, dist)'''
    copy = ranges_valid.copy()
    copy[copy == 0] = 16
    argmin = np.argmin(copy)

    angle = argmin * RESOLUTION
    dist = copy[argmin]
    
    return np.array([angle, dist, argmin])


def find_circle(ranges):
    max_d = np.max(ranges)
    thre = max_d / 2

    bin = np.zeros_like(ranges, dtype=np.uint8)
    bin[np.where(ranges > thre)] = 1

    median = bin.copy()    # 中值滤波

    length = len(bin)
    for i in range(1, length - 1):
        if bin[i - 1] == 1 and bin[i + 1] == 1:
            median[i] = 1
    
    # ---|---|------------------------|----|----
    # - l0 - r0 -------------------- l1 -- r1 --

    status = 0
    l0 = r0 = 0     # 呼啦圈两个顶点的左右边界
    l1 = r1 = 0
    for i in range(length - 1):
        if status == 0:
            if median[i] == 1:
                l0 = i
                status = 1

        elif status == 1:
            if median[i] == 1 and median[i + 1] == 0:
                r0 = i
                status = 2

        elif status == 2:
            if median[i] == 1:
                l1 = i
                status = 3

        elif status == 3:
            if median[i] == 1 and median[i + 1] == 0:
                r1 = i
                break

    if not l0 < r0 or not r0 < l1 or not l1 < r1:
        return None
    
    if l0 == 0 or r1 == len(median) - 1:
        return None
    
    if r0 - l0 < 2 or r0 - l0 > 25: # 两个边缘
        return None
    
    if r1 - l1 < 2 or r1 - l1 > 25: # 中间间距
        return None
    
    if l1 - r0 < 60:
        return None
    
    l = (l0 + r0) // 2
    r = (l1 + r1) // 2
    ld = ranges[l]
    rd = ranges[r]

    if ld == 0 or rd == 0:
       return None

    mid0 = map2xy(l, ld)
    mid1 = map2xy(r, rd)
    dist = calc_dist(mid0 ,mid1)
    
    print(l0, r0, l1, r1)
    print(dist, l, r, ld, rd)

# 相关运算，并可消除位移
def relate(arr, kernel):
    length = len(arr)
    start = - int(len(kernel)/2)
    res = [] 
    for i in range(start, length + start):
        value = 0
        for k in range(len(kernel)):
            id = i + k
            if id >= length:
                id -= length 
            value += arr[id] * kernel[k]
        res.append(abs(value))   
    return res


k1 = [-1, 1]
def find_circle2(ranges):
    grad = np.array(relate(ranges, k1))
    thre = np.max(grad) / 2.5

    bin = np.zeros_like(grad)
    bin[np.where(grad > thre)] = 1

    length = len(grad)
    rising = []
    res = []

    for i in range(length - 1):
        if bin[i] == 0 and bin[i + 1] == 1: # 上升沿
            rising.append(i)
        
        if len(rising) == 2:
            res.append(( rising[0] + rising[1]) // 2 )
            rising = []

        if len(res) == 2:
            break
    
    if len(res) != 2:
        return None
    
    if res[1] - res[0] < 20 or res[1] - res[0] > 300:
        return None
    
    if ranges[res[0]] == 0 or ranges[res[1]] == 0:
        return None
    
    p0 = map2xy(res[0], ranges[res[0]]) 
    p1 = map2xy(res[1], ranges[res[1]])

    return np.array((p0, p1))

def find_circle3(ranges, circle_diameter = 110, log = False):
    grad = np.array(relate(ranges, k1))
    thre = np.max(grad) / 3
    bin = np.zeros_like(grad)
    bin[np.where(grad > thre)] = 1

    # print(ranges)
    # print(grad)
    # print(bin)


    length = len(grad)
    rising = []
    res = []

    for i in range(length - 1):
        if bin[i] == 0 and bin[i + 1] == 1: # 上升沿
            rising.append(i)
        
        if len(rising) == 2:
            edge_index_0 = _find_nearest_valid(ranges, rising[0])   # avoid offset
            edge_index_1 = _find_nearest_valid(ranges, rising[1])
            if edge_index_0 is None or edge_index_1 is None:
                rising.clear()
                continue

            p0 = map2xy(edge_index_0, ranges[edge_index_0])
            p1 = map2xy(edge_index_1, ranges[edge_index_1])
            d = calc_dist(p0, p1)

            # print("find edge", p0, p1, d)
            if d > 0.5 and d < 5:
                res.append( ( edge_index_0 + edge_index_1) // 2 )
                rising.clear()
            else:
                rising.pop(0)   # 从第二个开始，重新开始匹配
        
        if len(res) == 2:
            break

    circle_pos = None

    if len(res) == 2:
        p0 = map2xy(res[0], ranges[res[0]]) 
        p1 = map2xy(res[1], ranges[res[1]])
            

        circle_pos = np.array([p0, p1])

    elif len(res) == 1:
        print("only 1 edge detected")

        # avoid finding the behind point
        index = res[0]
        nearest_index = index
        for i in range(-20, 20 + 1):
            if ranges[index + i] > 30 and ranges[index + i] < ranges[nearest_index]:
                nearest_index = index + i

        p0 = map2xy(nearest_index, ranges[nearest_index]) 
        d = calc_dist([0., 0.], p0)
        p1 = map2xy(nearest_index, d + circle_diameter)
            

        circle_pos = np.array([p0, p1])  


    if log:
        try:
            with open("/home/ninetripod/catkin_ws/src/uav/log/radar.json", "a+") as f:
                json.dump({
                    "pos": circle_pos.tolist(),
                    "index": res,
                    "range": ranges.tolist()
                }, f)
                f.write("\n\n")

        except Exception as e:
            print(e)

    return circle_pos


def find_circle4(ranges, circle_diameter = 110, log = False):
    '''将雷达数据二值化；根据一组上下降沿寻找边缘，并判断长度；直到找到两个边缘或遍历完毕为止'''
    thre = np.max(ranges) / 3

    bin = np.zeros_like(ranges)
    bin[np.where(ranges > thre)] = 1

    length = len(ranges)

    pair = []
    edges = []
    status = 0

    for i in range(length - 1):
        if status == 0: # 找上升沿
            if bin[i] == 0 and bin[i + 1] == 1:
                pair.append(i + 1)
                status = 1
        elif status == 1: # 找下降沿
            if bin[i] == 1 and bin[i + 1] == 0:
                pair.append(i)
                status = 0

        if len(pair) == 2:
            # 找到边缘
            p0 = map2xy(pair[0], ranges[pair[0]])
            p1 = map2xy(pair[1], ranges[pair[1]])
            d = calc_dist(p0, p1)

            # 解决太远，只扫到一个的情况
            if pair[0] == pair[1]:
                d += RESOLUTION * ranges[pair[0]]

            print("find pair", pair, ranges[pair[0]: pair[1] + 1], d)

            if d > 0.4 and d < 10 or d > 80:
                edges.append(( pair[0] + pair[1]) // 2 )

            pair.clear()
        
        if len(edges) == 2:
            # 找到了两条边
            break
    
    circle_pos = None
    
    if len(edges) == 2:
        p0 = map2xy(edges[0], ranges[edges[0]])
        p1 = map2xy(edges[1], ranges[edges[1]])

        circle_pos = np.array([p0, p1])


    elif len(edges) == 1:
        print("only 1 edge detected")
        
        index = edges[0]

        # avoid finding the behind point
        index = edges[0]
        nearest_index = index

        for i in range(-20, 20 + 1):
            if ranges[index + i] > 30 and ranges[index + i] < ranges[nearest_index]:
                nearest_index = index + i

        p0 = map2xy(nearest_index, ranges[nearest_index]) 
        d = calc_dist([0., 0.], p0)
        p1 = map2xy(nearest_index, d + circle_diameter)

        print("nearest", (nearest_index, ranges[nearest_index]))

        circle_pos = np.array([p0, p1])

    
    if log:
        try:
            with open("/home/ninetripod/catkin_ws/src/uav/log/radar.json", "a+") as f:
                json.dump({
                    "pos": circle_pos.tolist(),
                    "index": circle_pos,
                    "range": ranges.tolist()
                }, f)
                f.write("\n\n")

        except Exception as e:
            print(e)


    return circle_pos




def _find_nearest_valid(ranges, index, size = 2):
    if index < size:
        return None
    
    for i in range(size + 1):
        if ranges[index + i] != 0:
            return index + i
        
        elif ranges[index - i] != 0:
            return index - i

    return None


map2xy = lambda index, dist: (np.cos(index * RESOLUTION) * dist ,  np.sin(index * RESOLUTION) * dist) 

calc_dist = lambda p0, p1: np.sqrt((p0[0] - p1[0])**2 + (p0[1] - p1[1])**2)

# slope = lambda p0, p1: 

# 将0到90度的数据，映射到300cm x 300cm的图片
def map2img(ranges, size = 300):
    res = np.zeros((size, size), dtype = np.uint8)
    
    length = len(ranges)
    for i in range(length):
        d = ranges[i]
        if d > size or d < 30 : continue

        x = np.cos(i * RESOLUTION) * d
        y = np.sin(i * RESOLUTION) * d
        res[int(x), int(y)] = 0xFF
    
    return res